import base64
from Crypto.Cipher import AES
import sys
import re


class AES128_UDS:
    def __init__(self, Learmask=None) -> None:
        if Learmask is None:
            self.key = 0x01020304050607080102030405060708.to_bytes(length=16)
        else:
            self.key = Learmask

        self.mode = AES.MODE_ECB

    def pad(self, Seed=None):
        '''
        尝试的pkcs5补齐方法
        '''

        if Seed is None:
            self.data = 0x6816b4d5.to_bytes(length=4)
        else:
            self.data = Seed
        count = len(self.data)

        add_count = 16 - count % 16
        pad_item = add_count
        add_data = pad_item.to_bytes(1) * add_count

        # print(f'[file:{__file__}]=>[line:{
        #     (sys._getframe().f_lineno)}]==>add_data{add_data}')
        dataseed = self.data + add_data
        # print(f'[file:{__file__}]=>[line:{
        #     (sys._getframe().f_lineno)}]==>data{dataseed}')

        return dataseed

    def AES_en(self, data):
        # 创建加密对象
        AES_obj = AES.new(self.key, self.mode)

        # 补齐加密数据
        pad_data = self.pad(data)
        # 完成加密
        AES_en_byte = AES_obj.encrypt(pad_data)

        # 转为hex字符串
        AES_en_str = AES_en_byte.hex()

        return AES_en_byte

    def AES_de(self, data):
        '''
            待实现
        '''
        pass


if 0:
    # 参考算法实现素材
    def AES_en(key, data):
        # 将长度不足16字节的字符串补齐
        if len(data) < 16:
            data = pad(data)
        # 创建加密对象
        # AES_obj = AES.new(key.encode("utf-8"), AES.MODE_CBC, iv.encode("utf-8"))
        AES_obj = AES.new(key.encode("utf-8"), AES.MODE_ECB)
        # 完成加密
        AES_en_str = AES_obj.encrypt(data)
        # 用base64编码一下
        AES_en_str = base64.b64encode(AES_en_str)
        # 最后将密文转化成字符串
        AES_en_str = AES_en_str.decode("utf-8")
        return AES_en_str

    def AES_de(key, data):
        # 解密过程逆着加密过程写
        # 将密文字符串重新编码成二进制形式
        data = data.encode("utf-8")
        # 将base64的编码解开
        data = base64.b64decode(data)
        # 创建解密对象
        # AES_de_obj = AES.new(key.encode("utf-8"), AES.MODE_CBC, iv.encode("utf-8"))
        AES_de_obj = AES.new(key.encode("utf-8"), AES.MODE_ECB)
        # 完成解密
        AES_de_str = AES_de_obj.decrypt(data)
        # 去掉补上的空格
        AES_de_str = AES_de_str.strip()
        # 对明文解码
        AES_de_str = AES_de_str.decode("utf-8")
        return AES_de_str

    key = 'miyaoxuyao16ziji'
    data = b'hello world'
    # 将原始的明文用空格填充到16字节
    data = AES_en(key, data)
    print(data)
    data = AES_de(key, data)
    print(data)

    class AESECB:
        '''
            AESECB算法示例类
            [ ] 该类填充可能也不是pkcs5，需要进一步确认
        '''

        from Crypto.Cipher import AES

        def __init__(self, key):
            self.key = key
            self.mode = AES.MODE_ECB
            self.bs = 16  # block size
            self.PADDING = lambda s: s + \
                (self.bs - len(s) % self.bs) * chr(self.bs - len(s) % self.bs)

        def encrypt(self, text):
            generator = AES.new(self.key, self.mode)  # ECB模式无需向量iv
            crypt = generator.encrypt(self.PADDING(text))
            crypted_str = base64.b64encode(crypt)
            result = crypted_str.decode()
            return result

        def decrypt(self, text):
            generator = AES.new(self.key, self.mode)  # ECB模式无需向量iv
            text += (len(text) % 4) * '='
            decrpyt_bytes = base64.b64decode(text)
            meg = generator.decrypt(decrpyt_bytes)
            # 去除解码后的非法字符
            try:
                result = re.compile(
                    '[\\x00-\\x08\\x0b-\\x0c\\x0e-\\x1f\n\r\t]').sub('', meg.decode())
            except Exception:
                result = '解码失败，请重试!'
            return result

    class EncryptDate2:
        '''
            备选第二个AEB类
            该类用的是pkcs7填充，不能直接使用
        '''

        import base64
        from Crypto.Cipher import AES
        from Crypto.Util.Padding import pad

        def __init__(self, key):
            # 初始化密钥
            self.key = key
            # 初始化数据块大小
            self.length = AES.block_size
            # 初始化AES,ECB模式的实例
            self.aes = AES.new(self.key.encode("utf-8"), AES.MODE_ECB)
            # 截断函数，去除填充的字符
            self.unpad = lambda date: date[0:-ord(date[-1])]

        def fill_method(self, aes_str):
            '''pkcs7补全'''
            pad_pkcs7 = pad(aes_str.encode('utf-8'),
                            AES.block_size, style='pkcs7')

            return pad_pkcs7

        def encrypt(self, encrData):
            # 加密函数,使用pkcs7补全
            res = self.aes.encrypt(self.fill_method(encrData))
            # 转换为base64
            msg = str(base64.b64encode(res), encoding="utf-8")

            return msg

        def decrypt(self, decrData):
            # base64解码
            res = base64.decodebytes(decrData.encode("utf-8"))
            # 解密函数
            msg = self.aes.decrypt(res).decode("utf-8")

            return self.unpad(msg)

    if 0:
        # key的长度需要补长(16倍数),补全方式根据情况而定,未补齐会报错
        # key字符长度决定加密结果,长度16：加密结果AES(128),长度32：结果就是AES(256)
        eg = EncryptDate2("asdfghjklqwertyu")
        # 加密字符串长同样需要16倍数：需注意,不过代码中pad()方法里，帮助实现了补全（补全方式就是pkcs7）
        en = eg.encrypt("测试数据")
        de = eg.decrypt(en)
        print(f"加密结果：{en}")
        print(f"解密结果：{de}")

    # pad1方法
    def pad1(data):
        pad_data = data
        for i in range(0, 16-len(data)):
            pad_data = pad_data + b' '
        return pad_data
